#include "SocketEUDP.h"

#include <time.h>
#pragma comment(lib, "Ws2_32.lib")
#pragma comment(lib, "winmm.lib")

#ifdef _DEBUG

#define ASSERT(x) assert(x)

#include <assert.h>
#include <iostream>

#else

#define ASSERT(x)

#endif

bool SocketEUDP::CreateSocket(int port)
{
	//Startup the winsock service
	WSADATA wsaData;
	if(WSAStartup(MAKEWORD(2, 2), &wsaData) != 0) 
	{
#ifdef _DEBUG
		std::cout << "WSAStartup failed." << std::endl;
#endif
		return false;
	}

	if( !CreateInboundSocket(port) )
	{
		return false;
	}

	if( !CreateOutboundSocket() )
	{
		return false;
	}

	return true;
}

bool SocketEUDP::CreateInboundSocket(int port)
{
	ASSERT(m_socketIn == INVALID_SOCKET);

	sockaddr_in service;
	service.sin_family = AF_INET;
	service.sin_addr.s_addr = 0; //Allow connections to this port from all ip's
	service.sin_port = htons((u_short)port);

	m_socketIn = socket(AF_INET, SOCK_DGRAM, IPPROTO_UDP);

	if(m_socketIn == SOCKET_ERROR)
	{
#if _DEBUG
		std::cout << "Error creating the socket." << WSAGetLastError() << std::endl;
#endif

		m_socketIn = INVALID_SOCKET;
		return false;
	}

	if( bind(m_socketIn, 
		     (SOCKADDR*) &service, 
			 sizeof(service)
			) == SOCKET_ERROR )
	{
#if _DEBUG
		std::cout << "bind() failed with error " << WSAGetLastError() << std::endl;
#endif

		closesocket(m_socketIn);
		m_socketIn = INVALID_SOCKET;
		return false;;
	}

	if (m_socketIn != INVALID_SOCKET)
	{
		//Successfully created an inbound socket, store the port
		m_port = (u_short)port;
#ifdef _DEBUG
		std::cout << "Bound socket to port " << port << std::endl;
#endif
		return true;
	}

#if _DEBUG
	std::cout << "Error creating and binding the socket" << std::endl;
#endif
	return false;
}

bool SocketEUDP::CreateOutboundSocket()
{
	//Attempt to create an outbound socket
	m_socketOut = socket(AF_INET, SOCK_DGRAM, IPPROTO_UDP);

	if(m_socketOut == SOCKET_ERROR || m_socketOut == INVALID_SOCKET)
	{
#ifdef _DEBUG
		std::cout << "Error creating and binding socket." << std::endl;
#endif

		m_socketOut = INVALID_SOCKET;
		return false;
	}
	if (m_socketOut != INVALID_SOCKET && m_socketOut != SOCKET_ERROR)
	{
		return true;
	}

#ifdef _DEBUG
	std::cout << "Unknown error creating and binding socket." << std::endl;
#endif

	m_socketOut = INVALID_SOCKET;
	return false;
}

EUDP_R SocketEUDP::ConnectTo(std::string ip, int port)
{
	//Attempt to send a connection message to the specified ip
	sockaddr_in addr;
	addr.sin_family = AF_INET;
	addr.sin_port = htons((u_short)port);
	addr.sin_addr.s_addr = inet_addr(ip.c_str());
	int flags = 0;

	//Flag the message as a connection message
	Message message;
	message.messageTypeEUDP = (EUDP_M)htons(EUDP_M_CONNECT);
	message.port = htons(m_port);

	//Create CRC
	message.checksum = CreateChecksum(message);

	if( sizeof(Message) == sendto( m_socketOut, (const char*)&message, sizeof(Message), flags, (SOCKADDR*)&addr, sizeof(addr) ) )
	{
		//Message sent successfully, store details in connection info if this is the establishing handshake
		int ref = ConnectionReference(ip, port);
		if( ref == -1 )
		{
			//Sending first handshake, add to the connection list
			ConnectionAdd(ip, port, 0);
			ConnectTo(ip, port);
#ifdef _DEBUG
			std::cout << "Sending first handshake to "<< ip << " on port " << port << std::endl;
#endif
		}

		return EUDP_R_OK;
	}

#ifdef _DEBUG
	std::cout << "Error sending connection message" << std::endl;
#endif

	return EUDP_R_ERROR;
}

EUDP_R SocketEUDP::Send(int reference, short messageType, void *data, size_t length, bool reliable)
{
	//Check that we have a connection with the specified ip
	if( ConnectionCheck(reference) != -1)
	{
		if( ConnectionCheck(reference) < 2 )
		{
			//Connection not yet established
			return EUDP_R_NO_HANDSHAKE;
		}
		
		sockaddr_in addr;
		addr.sin_family = AF_INET;
		addr.sin_port = htons(m_connections[reference].port);
		addr.sin_addr.s_addr = inet_addr(m_connections[reference].ip.c_str());
		int flags = 0;

		//Construct the message to send
		Message message;
		message.messageType = htons(messageType);
		message.port        = htons(m_port);
		message.length      = htons((u_short)length);
		char *cData = static_cast<char*>(data);
		std::copy(cData, cData + length, message.data);

		if(reliable)
		{
			message.messageTypeEUDP = (EUDP_M)htons(EUDP_M_MESSAGE_RELIABLE);
			message.reference = m_connections[reference].reliableReference;
			++m_connections[reference].reliableReference;

			ReliableInfo info;
			info.connectionReference = reference;
			info.reliableReference = message.reference;
			m_messagesToSend.insert( std::pair<ReliableInfo, Message>(info, message) );
			m_messagesToSendTimer.insert( std::pair<ReliableInfo, DWORD>(info, timeGetTime()) );

		} else {
			message.messageTypeEUDP = (EUDP_M)htons(EUDP_M_MESSAGE);
			message.reference = 0;
		}

		//Create CRC
		message.checksum = CreateChecksum(message);

		if(sizeof(Message) == sendto( m_socketOut, (const char*)&message, sizeof(Message), flags, (SOCKADDR*)&addr, sizeof(addr) ) )
		{
			return EUDP_R_OK;
		}

		return EUDP_R_ERROR;
	}

	//No connection with the specified ip
	return EUDP_R_NO_CONNECTION;
}

EUDP_R SocketEUDP::Disconnect(int reference)
{
	//TODO INCOMPLETE
	reference += 1;
	return EUDP_R_ERROR;
}


int SocketEUDP::GetMessageCount()
{
	return m_messagesToReceive.size();
}

SocketEUDP::Message SocketEUDP::GetMessages()
{
	Message returnMessage;
	
	if(m_messagesToReceive.size() > 0)
	{
		returnMessage = m_messagesToReceive.front();
		m_messagesToReceive.pop();
		return returnMessage;
	}

	//No messages waiting to be collected, return an empty message
	returnMessage.messageTypeEUDP = EUDP_M_EMPTY;
	return returnMessage;
}

void SocketEUDP::Update()
{
	UpdateOutbound();
	UpdateInbound();
}

void SocketEUDP::UpdateInbound()
{
	//Check if there are any messages waiting to come in
	timeval timeout;
	timeout.tv_sec = 0;
	timeout.tv_usec = 0;

	fd_set sockets;
	FD_ZERO(&sockets);
	FD_SET(m_socketIn, &sockets);

	while( select(0, &sockets, NULL, NULL, &timeout) != 0 )
	{
		//Message waiting, receive message
		char data[2048];
		int dataLength;
		std::string ip;

		int flags = 0;
		sockaddr_in fromAddressIP;
		int fromLength = sizeof(fromAddressIP);
		dataLength = recvfrom(m_socketIn, data, BUFFER_MAX, flags, (SOCKADDR*)&fromAddressIP, &fromLength);
		
		ip = inet_ntoa(fromAddressIP.sin_addr);
		
		//Convert data into a message and change from network to host format
		Message message         = *(Message*)data;
		message.messageTypeEUDP = (EUDP_M)ntohs((u_short)message.messageTypeEUDP);
		message.reference       = ntohl(message.reference);
		message.port            = ntohs(message.port);
		message.messageType     = ntohs(message.messageType);
		message.length          = ntohs(message.length);

		//Validate the CRC
		//TODO: CRC disabled, unable to validate when data is present
		if( true )
		{
			//Check if we have a connection with the sender
			int ref = ConnectionReference(ip, message.port);
			if( ref == -1 )
			{
				//No connection, check if it is a connection request
				if(message.messageTypeEUDP == EUDP_M_CONNECT)
				{
					if(m_autoConnect)
					{
						//Add to the connection list and return the second handshake
						ConnectionAdd(ip, message.port, 1);
						ConnectTo(ip, message.port);
#ifdef _DEBUG
						std::cout << "Sending second handshake to "<< ip << " on port " << message.port << std::endl;
#endif
						break;

					} else {
						//TODO: Handling for not automatically connecting
					}
				}
			}

			//Check that we have completed handshaking
			if( ( ConnectionCheck(ref) != -1 ) && ( ConnectionCheck(ref) < 2 ) )
			{
				//Haven't finished the handshake, treat this message as a connection message and increase the shake count by two
				m_connections[ref].shakeCount += 2;
				m_connections[ref].connectionTime = timeGetTime();
				m_connections[ref].lastMessageTime = timeGetTime();
				//If this is the second handshake, send the last handshake message
				if(m_connections[ref].shakeCount == 2)
				{
					ConnectTo(ip, message.port);
#ifdef _DEBUG
					std::cout << "Sending third handshake to "<< ip << " on port " << message.port << std::endl;
#endif
				}

				//Formed a connection on this end
#ifdef _DEBUG
				std::cout << "Connection Established" << std::endl;
#endif
				m_events.push("Connected To: " + m_connections[ref].ip);
			
			} else {
			
				//We are connected to the sender, check the message type
				switch(message.messageTypeEUDP)
				{
				case EUDP_M_CONNECT:
					//Received a connection message, do nothing
					break;

				case EUDP_M_SYNC:
					//Received a sync message, set lastMessageTime to prevent timeout
					m_connections[ref].lastMessageTime = timeGetTime();
#ifdef _DEBUG
					std::cout << "Received Sync Message" << std::endl;
#endif
					break;

				case EUDP_M_MESSAGE:
					//A regular message, put into the message queue
					message.reference = ref;
					m_messagesToReceive.push(message);
					break;

				case EUDP_M_MESSAGE_RELIABLE:
					//Received a reliable message, send back a confirmation message
					SendReliableConfirmation(ref, message.reference);
					message.reference = ref;
					m_messagesToReceive.push(message);
					break;

				case EUDP_M_MESSAGE_RELIABLE_CONFIRM:
					//Received a reliable message confirmation, remove the message from the resend list
					ReliableInfo info;
					info.connectionReference = ref;
					info.reliableReference = message.reference;
					m_messagesToSend.erase(info);
					m_messagesToSendTimer.erase(info);
					break;

				default:
#ifdef _DEBUG
					std::cout << "Message received with invalid EUDP message type" << std::endl;
#endif
					break;
				}
			}
		}
	}
	
}

void SocketEUDP::UpdateOutbound()
{
	//Check if there are any messages in the reliable message map that need sending
	if(m_messagesToSend.size() != 0)
	{
		//If the time since the message has been sent has exceeded a certain limit, resend the message
		for(std::map<ReliableInfo, DWORD>::iterator i = m_messagesToSendTimer.begin(); i != m_messagesToSendTimer.end(); ++i)
		{
			if( i->second < (timeGetTime() - RELIABLE_RESEND) )
			{
				sockaddr_in addr;
				addr.sin_family = AF_INET;
				addr.sin_port = htons(m_connections[i->first.connectionReference].port);
				addr.sin_addr.s_addr = inet_addr(m_connections[i->first.connectionReference].ip.c_str());
				int flags = 0;


				m_events.push("Resending Reliable Message");
				if(sizeof(Message) == sendto( m_socketOut, (const char*)&m_messagesToSend[i->first], sizeof(Message), flags, (SOCKADDR*)&addr, sizeof(addr) ) )
				{
					//Message sent, set last send timer
					i->second = timeGetTime();
				}
			}
		}
	}

	
	
	//Check through each connection, check if it has timed out and if a sync message needs to be sent
	std::vector<int> deleteConnections;

	for(std::map<int, ConnectionInfo>::iterator i = m_connections.begin(); i != m_connections.end(); ++i)
	{
		//Check if it a sync message is required
		if( (i->second.lastMessageTime) < (timeGetTime() - SYNC) )
		{
			//First check if the connection has timed out
			if( (i->second.lastMessageTime) < (timeGetTime() - TIMEOUT) )
			{
				//TIMEOUT
				deleteConnections.push_back(i->first);				

			//Not timed out, check if a sync message needs to be sent
			} else if( (i->second.lastSyncMessageTime) < (timeGetTime() - SYNC) ) {
				//Send a sync message
				sockaddr_in addr;
				addr.sin_family = AF_INET;
				addr.sin_port = htons(i->second.port);
				addr.sin_addr.s_addr = inet_addr(i->second.ip.c_str());
				int flags = 0;

				//Construct the message to send
				Message message;
				message.messageType = 0;
				message.port        = htons(m_port);
				message.length      = 0;
				message.messageTypeEUDP = (EUDP_M)htons(EUDP_M_SYNC);

				if(sizeof(Message) == sendto( m_socketOut, (const char*)&message, sizeof(Message), flags, (SOCKADDR*)&addr, sizeof(addr) ) )
				{
					//Message sent successfully, update last sync message time
					i->second.lastSyncMessageTime = timeGetTime();
				} else {

#ifdef _DEBUG
					std::cout << "Error sending sync message!" << std::endl;
#endif

				}
			}
		}
	}

	//Delete any connections that have been marked
	for(unsigned int i = 0; i < deleteConnections.size(); ++i)
	{
		m_events.push("Connection: " + m_connections[i].ip + " Timed Out");
		
		ConnectionRemove(deleteConnections[i]);
		
#ifdef _DEBUG
		std::cout << "CONNECTION TIMED OUT" << std::endl;
#endif
	}
}

EUDP_R SocketEUDP::SendReliableConfirmation(int reference, unsigned long reliableReference)
{
	//Send a confirmation message
	sockaddr_in addr;
	addr.sin_family = AF_INET;
	addr.sin_port = htons(m_connections[reference].port);
	addr.sin_addr.s_addr = inet_addr(m_connections[reference].ip.c_str());
	int flags = 0;

	//Flag the message as a connection message
	Message message;
	message.messageTypeEUDP = (EUDP_M)htons(EUDP_M_MESSAGE_RELIABLE_CONFIRM);
	message.port = htons(m_port);
	message.reference = htonl(reliableReference);

	m_events.push("Sending Reliable Message Confirmation");

	if( sizeof(Message) == sendto( m_socketOut, (const char*)&message, sizeof(Message), flags, (SOCKADDR*)&addr, sizeof(addr) ) )
	{
		//Message sent successfully, store details in connection info if this is the establishing handshake
		return EUDP_R_OK;
	}

#ifdef _DEBUG
	std::cout << "Error sending message confirmation" << std::endl;
#endif

	return EUDP_R_ERROR;
}

void SocketEUDP::ConnectionAdd(std::string ip, int port, int shakeCount)
{
	//Add a new connection to the connection list
	ConnectionInfo newConnection;
	newConnection.ip = ip;
	newConnection.port = (u_short)port;
	newConnection.shakeCount = shakeCount;

	newConnection.connectionTime = 0;
	newConnection.lastMessageTime = timeGetTime();

	m_connections.insert(std::pair<int, ConnectionInfo>(m_connectionMapValue, newConnection));
	++m_connectionMapValue;
}

int SocketEUDP::ConnectionCheck(int reference)
{
	for(std::map<int, ConnectionInfo>::iterator i = m_connections.begin(); i != m_connections.end(); ++i)
	{
		if( reference == i->first )
		{
			return i->second.shakeCount;
		}
	}

	return -1;
}

int SocketEUDP::ConnectionReference(std::string ip, int port)
{
	for(std::map<int, ConnectionInfo>::iterator i = m_connections.begin(); i != m_connections.end(); ++i)
	{
		if( (port == i->second.port) && (ip == i->second.ip) )
		{
			return i->first;
		}
	}

	return -1;
}

void SocketEUDP::ConnectionRemove(int reference)
{
	m_connections.erase(reference);
}

bool SocketEUDP::ConnectedTEMP()
{
	//Yeah, this is temporary
	if(ConnectionCheck(0) >= 2)
	{
		return true;
	}
	return false;
}

bool operator < (const SocketEUDP::ReliableInfo a, const SocketEUDP::ReliableInfo b)
{
	float first =  ( a.reliableReference + (1.0f / a.connectionReference) ); 
	float second = ( b.reliableReference + (1.0f / b.connectionReference) ); 
	if(first < second)
	{
		return true;
	}

	return false;
}

char SocketEUDP::CreateChecksum(Message message)
{
	char returnChecksum;

	for(char *i = ((char*)&message + 1); i < ((char*)&message + MESSAGE_SIZE(message.length)); ++i)
	{
		returnChecksum ^= *i;
	}

	return returnChecksum;
}

bool SocketEUDP::ValidateChecksum(Message message)
{
	char checksum;
	for(char *i = ((char*)&message + 1); i < ((char*)&message + MESSAGE_SIZE(message.length)); ++i)
	{
		checksum ^= *i;
	}

	if(checksum == message.checksum)
	{
		return true;
	}

	return false;
}

std::string SocketEUDP::GetLastEvent()
{
	std::string returnString;
	if(m_events.empty() == false)
	{
		returnString = m_events.front();
		m_events.pop();
	} else {
		returnString = "No Events";
	}

	return returnString;
}

std::map<int, SocketEUDP::ConnectionInfo> SocketEUDP::GetConnections()
{
	return m_connections;
}